#include "DeferredProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Game/GameState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "../LocalShader.h"
#include "../LocalScene.h"
#include "../ResourceProcess.h"
#include "../Shadows/ShadowsAtrousFilterProcess.h"
#include "../Reflections/ReflectionsAtrousFilterProcess.h"

void DeferredProcess::OnInit()
{
	mLocalScene = static_cast<LocalScene*>(mGameScene);
	mResourceProcess = mLocalScene->GetProcess<ResourceProcess>("ResourceProcess");
	mShadowsAtrousFilterProcess = mLocalScene->GetProcess<ShadowsAtrousFilterProcess>("ShadowsAtrousFilterProcess");
	mReflectionsAtrousFilterProcess = mLocalScene->GetProcess<ReflectionsAtrousFilterProcess>("ReflectionsAtrousFilterProcess");


	mShadow = true;
	mAo = true;
	mGi = false;
	mReflections = true;

	mPushConstants.shadow = static_cast<int32>(mShadow);
	mPushConstants.ao = static_cast<int32>(mAo);
	mPushConstants.gi = static_cast<int32>(mGi);
	mPushConstants.reflections = static_cast<int32>(mReflections);

	ENQUEUE_RENDER_COMMAND(BuildModel)([this](VkCommandBuffer commandBuffer)
		{
			RefCountPtr<RHISampler> trilinearSampler = GetSampler("ResourceProcess", "trilinearSampler");
			RefCountPtr<RHISampler> bilinearSampler = GetSampler("ResourceProcess", "bilinearSampler");
			RefCountPtr<RHISampler> nearestSampler = GetSampler("ResourceProcess", "nearestSampler");

			std::array<RefCountPtr<RHITexture>, 2> gBuffer1;
			std::array<RefCountPtr<RHITexture>, 2> gBuffer2;
			std::array<RefCountPtr<RHITexture>, 2> gBuffer3;
			std::array<RefCountPtr<RHITexture>, 2> gBufferDepth;

			gBuffer1[0] = GetTexture("GBufferProcess", "gBuffer1_s");
			gBuffer2[0] = GetTexture("GBufferProcess", "gBuffer2_s");
			gBuffer3[0] = GetTexture("GBufferProcess", "gBuffer3_s");
			gBufferDepth[0] = GetTexture("GBufferProcess", "gBufferDepth_s");

			gBuffer1[1] = GetTexture("GBufferProcess", "gBuffer1_d");
			gBuffer2[1] = GetTexture("GBufferProcess", "gBuffer2_d");
			gBuffer3[1] = GetTexture("GBufferProcess", "gBuffer3_d");
			gBufferDepth[1] = GetTexture("GBufferProcess", "gBufferDepth_d");

			std::array<RefCountPtr<RHITexture>, 2> shadowImage;
			shadowImage[0] = GetTexture("ShadowsAtrousFilterProcess", "outputImage_s");
			shadowImage[1] = GetTexture("ShadowsAtrousFilterProcess", "outputImage_d");

			RefCountPtr<RHITexture> aoImage = GetTexture("AoDenoiseBilateralBlurProcess", "outputImage_d");

			RefCountPtr<RHITexture> giImage = GetTexture("DdgiSampleProbeGridProcess", "outputImage");

			std::array<RefCountPtr<RHITexture>, 2> reflectionsImage;
			reflectionsImage[0] = GetTexture("ReflectionsAtrousFilterProcess", "outputImage_s");
			reflectionsImage[1] = GetTexture("ReflectionsAtrousFilterProcess", "outputImage_d");


			RefCountPtr<RHIBuffer> perFrameUbo = GetBuffer("ResourceProcess", "perFrameUbo");
			RefCountPtr<RHITexture> shImage = GetTexture("SphericalHarmonicsProcess", "shImage");
			RefCountPtr<RHITexture> prefilterImage = GetTexture("PrefilterProcess", "prefilterImage");
			RefCountPtr<RHISampler> prefilterSampler = GetSampler("PrefilterProcess", "prefilterSampler");
			RefCountPtr<RHITexture> brdfLutImage = GetTexture("PrefilterProcess", "brdfLutImage");



			mDescriptor = RHIDescriptor::Builder()
				.PushShader<DeferredVertexShader>()
				.PushShader<DeferredFragmentShader>()
				.SetDescriptorCount(0, 2) // GBuffer
				.SetDescriptorCount(1, 2) // Shadow
				.SetDescriptorCount(2, 1) // AO
				.SetDescriptorCount(3, 1) // GI
				.SetDescriptorCount(4, 2) // Reflect
				.SetDescriptorCount(5, 1) // Common
				.Build();


			{
				mDescriptor->SetValue("gBuffer1", gBuffer1[0], trilinearSampler);
				mDescriptor->SetValue("gBuffer2", gBuffer2[0], trilinearSampler);
				mDescriptor->SetValue("gBuffer3", gBuffer3[0], trilinearSampler);
				mDescriptor->SetValue("gBufferDepth", gBufferDepth[0], trilinearSampler);
				mGBufferDescriptorMask[0] = mDescriptor->AllocateDescriptor();

				mDescriptor->SetValue("gBuffer1", gBuffer1[1], trilinearSampler);
				mDescriptor->SetValue("gBuffer2", gBuffer2[1], trilinearSampler);
				mDescriptor->SetValue("gBuffer3", gBuffer3[1], trilinearSampler);
				mDescriptor->SetValue("gBufferDepth", gBufferDepth[1], trilinearSampler);
				mGBufferDescriptorMask[1] = mDescriptor->AllocateDescriptor();
			}

			{
				mDescriptor->SetValue("shadowImage", shadowImage[0], trilinearSampler);
				mShadowDescriptorMask[0] = mDescriptor->AllocateDescriptor();

				mDescriptor->SetValue("shadowImage", shadowImage[1], trilinearSampler);
				mShadowDescriptorMask[1] = mDescriptor->AllocateDescriptor();
			}

			{
				mDescriptor->SetValue("aoImage", aoImage, trilinearSampler);
				mAoDescriptorMask = mDescriptor->AllocateDescriptor();
			}

			{
				mDescriptor->SetValue("giImage", giImage, trilinearSampler);
				mGiDescriptorMask = mDescriptor->AllocateDescriptor();
			}

			{
				mDescriptor->SetValue("reflectionsImage", reflectionsImage[0], trilinearSampler);
				mReflectionsDescriptorMask[0] = mDescriptor->AllocateDescriptor();

				mDescriptor->SetValue("reflectionsImage", reflectionsImage[1], trilinearSampler);
				mReflectionsDescriptorMask[1] = mDescriptor->AllocateDescriptor();
			}

			{
				mDescriptor->SetValue("perFrameUBO", perFrameUbo);
				mDescriptor->SetValue("shImage", shImage, trilinearSampler);
				mDescriptor->SetValue("prefilterImage", prefilterImage, prefilterSampler);
				mDescriptor->SetValue("brdfLutImage", brdfLutImage, trilinearSampler);
				mCommonDescriptorMask = mDescriptor->AllocateDescriptor();
			}


			mDeferredImage = ColorRenderTarget::Builder()
				.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "deferredImage");

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(mDeferredImage)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mDeferredImage)
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build();

			mFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				.Build();
		});
}

void DeferredProcess::OnTick(float deltaTime)
{
	if (GGameState->GetFrameAmount() < 2)
		return;

	mPushConstants.shadow = static_cast<int32>(mShadow);
	mPushConstants.ao = static_cast<int32>(mAo);
	mPushConstants.gi = static_cast<int32>(mGi);
	mPushConstants.reflections = static_cast<int32>(mReflections);

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this](VkCommandBuffer commandBuffer)
		{

			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());


			VkDescriptorSet descriptorSets[] = {
				mGBufferDescriptorMask[mResourceProcess->mCurrMask].GetDescriptorSet(),
				mShadowDescriptorMask[mShadowsAtrousFilterProcess->mReadIndex].GetDescriptorSet(),
				mAoDescriptorMask.GetDescriptorSet(),
				mGiDescriptorMask.GetDescriptorSet(),
				mReflectionsDescriptorMask[mReflectionsAtrousFilterProcess->mReadIndex].GetDescriptorSet(),
				mCommonDescriptorMask.GetDescriptorSet()
			};


			vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 6, descriptorSets, 0, nullptr);


			vkCmdPushConstants(commandBuffer, mPipeline->GetPipelineLayout(), VK_SHADER_STAGE_FRAGMENT_BIT, 0, sizeof(PushConstants), &mPushConstants);


			vkCmdDraw(commandBuffer, 3, 1, 0, 0);

			vkCmdEndRenderPass(commandBuffer);
		});
}

void DeferredProcess::OnResize()
{
	RefCountPtr<RHISampler> trilinearSampler = GetSampler("ResourceProcess", "trilinearSampler");
	RefCountPtr<RHISampler> bilinearSampler = GetSampler("ResourceProcess", "bilinearSampler");
	RefCountPtr<RHISampler> nearestSampler = GetSampler("ResourceProcess", "nearestSampler");

	std::array<RefCountPtr<RHITexture>, 2> gBuffer1;
	std::array<RefCountPtr<RHITexture>, 2> gBuffer2;
	std::array<RefCountPtr<RHITexture>, 2> gBuffer3;
	std::array<RefCountPtr<RHITexture>, 2> gBufferDepth;

	gBuffer1[0] = GetTexture("GBufferProcess", "gBuffer1_s");
	gBuffer2[0] = GetTexture("GBufferProcess", "gBuffer2_s");
	gBuffer3[0] = GetTexture("GBufferProcess", "gBuffer3_s");
	gBufferDepth[0] = GetTexture("GBufferProcess", "gBufferDepth_s");

	gBuffer1[1] = GetTexture("GBufferProcess", "gBuffer1_d");
	gBuffer2[1] = GetTexture("GBufferProcess", "gBuffer2_d");
	gBuffer3[1] = GetTexture("GBufferProcess", "gBuffer3_d");
	gBufferDepth[1] = GetTexture("GBufferProcess", "gBufferDepth_d");

	std::array<RefCountPtr<RHITexture>, 2> shadowImage;
	shadowImage[0] = GetTexture("ShadowsAtrousFilterProcess", "outputImage_s");
	shadowImage[1] = GetTexture("ShadowsAtrousFilterProcess", "outputImage_d");

	RefCountPtr<RHITexture> aoImage = GetTexture("AoDenoiseBilateralBlurProcess", "outputImage_d");

	RefCountPtr<RHITexture> giImage = GetTexture("DdgiSampleProbeGridProcess", "outputImage");

	std::array<RefCountPtr<RHITexture>, 2> reflectionsImage;
	reflectionsImage[0] = GetTexture("ReflectionsAtrousFilterProcess", "outputImage_s");
	reflectionsImage[1] = GetTexture("ReflectionsAtrousFilterProcess", "outputImage_d");


	RefCountPtr<RHIBuffer> perFrameUbo = GetBuffer("ResourceProcess", "perFrameUbo");
	RefCountPtr<RHITexture> shImage = GetTexture("SphericalHarmonicsProcess", "shImage");
	RefCountPtr<RHITexture> prefilterImage = GetTexture("PrefilterProcess", "prefilterImage");
	RefCountPtr<RHISampler> prefilterSampler = GetSampler("PrefilterProcess", "prefilterSampler");
	RefCountPtr<RHITexture> brdfLutImage = GetTexture("PrefilterProcess", "brdfLutImage");

	{
		mDescriptor->SetValue("gBuffer1", gBuffer1[0], trilinearSampler);
		mDescriptor->SetValue("gBuffer2", gBuffer2[0], trilinearSampler);
		mDescriptor->SetValue("gBuffer3", gBuffer3[0], trilinearSampler);
		mDescriptor->SetValue("gBufferDepth", gBufferDepth[0], trilinearSampler);
		mDescriptor->UpdateDescriptor(mGBufferDescriptorMask[0]);

		mDescriptor->SetValue("gBuffer1", gBuffer1[1], trilinearSampler);
		mDescriptor->SetValue("gBuffer2", gBuffer2[1], trilinearSampler);
		mDescriptor->SetValue("gBuffer3", gBuffer3[1], trilinearSampler);
		mDescriptor->SetValue("gBufferDepth", gBufferDepth[1], trilinearSampler);
		mDescriptor->UpdateDescriptor(mGBufferDescriptorMask[1]);
	}

	{
		mDescriptor->SetValue("shadowImage", shadowImage[0], trilinearSampler);
		mDescriptor->UpdateDescriptor(mShadowDescriptorMask[0]);

		mDescriptor->SetValue("shadowImage", shadowImage[1], trilinearSampler);
		mDescriptor->UpdateDescriptor(mShadowDescriptorMask[1]);
	}

	{
		mDescriptor->SetValue("aoImage", aoImage, trilinearSampler);
		mDescriptor->UpdateDescriptor(mAoDescriptorMask);
	}

	{
		mDescriptor->SetValue("giImage", giImage, trilinearSampler);
		mDescriptor->UpdateDescriptor(mGiDescriptorMask);
	}

	{
		mDescriptor->SetValue("reflectionsImage", reflectionsImage[0], trilinearSampler);
		mDescriptor->UpdateDescriptor(mReflectionsDescriptorMask[0]);

		mDescriptor->SetValue("reflectionsImage", reflectionsImage[1], trilinearSampler);
		mDescriptor->UpdateDescriptor(mReflectionsDescriptorMask[1]);
	}

	{
		mDescriptor->SetValue("perFrameUBO", perFrameUbo);
		mDescriptor->SetValue("shImage", shImage, trilinearSampler);
		mDescriptor->SetValue("prefilterImage", prefilterImage, prefilterSampler);
		mDescriptor->SetValue("brdfLutImage", brdfLutImage, trilinearSampler);
		mDescriptor->UpdateDescriptor(mCommonDescriptorMask);
	}

	mDeferredImage = ColorRenderTarget::Builder()
		.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "deferredImage");

	mFrameBuffer = RHIFrameBuffer::Builder()
		.AddRenderTarget(mDeferredImage)
		.SetRenderPass(mRenderPass)
		.Build();
}
